/*
 * This file is part of Catfish.
 * Copyright (C) 2010 Namazu Studios LLC
 * 
 * Catfish is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as 
 * published by the Free Software Foundation, either version 3 of 
 * the License, or (at your option) any later version.
 * 
 * Catfish is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 * 
 * You should have received a copy of the GNU Lesser General Public
 * License along with Catfish.  If not, please visit:
 *  
 * http://www.gnu.org/licenses/
 *  
 */
package com.namazustudios.catfish.datasource;

import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import com.google.appengine.api.datastore.Entity;
import com.google.appengine.api.datastore.EntityNotFoundException;
import com.google.appengine.api.datastore.Key;
import com.google.appengine.api.datastore.Transaction;
import com.namazustudios.catfish.Catfish;
import com.namazustudios.catfish.DataSource;
import com.namazustudios.catfish.StackedDataSource;

/**
 * A DataSource for prefetching. When pushed on the stack, this datasource will
 * read all entities from its internal cache before deferring to any other data
 * source in the chain.
 * 
 * This is particularly useful if you know ahead of time the keys to all the
 * entities required to process a particular request.
 * 
 * @author patricktwohig
 * 
 */
public class PrefetchDataSource implements StackedDataSource {

	private Catfish catfish;

	private DataSource lower;

	private final Map<Key, Entity> cache;

	public PrefetchDataSource() {
		this.cache = new HashMap<Key, Entity>();
	}

	public PrefetchDataSource(Collection<Entity> entities) {

		cache = new HashMap<Key, Entity>();

		for (Entity entity : entities) {
			Key key = entity.getKey();

			if (key == null) {
				throw new IllegalArgumentException(
						"All entities must have a key.");
			} else {
				cache.put(key, entity);
			}
		}

	}

	@Override
	public void delete(Iterable<Key> keys) {
		getLower().delete(keys);
		for (Key key : keys) {
			cache.remove(key);
		}
	}

	@Override
	public void delete(Key... keys) {
		getLower().delete(keys);
		for (Key key : keys) {
			cache.remove(key);
		}
	}

	@Override
	public void delete(Transaction txn, Iterable<Key> keys) {
		getLower().delete(txn, keys);

		for (Key key : keys) {
			cache.remove(key);
		}

	}

	@Override
	public void delete(Transaction txn, Key... keys) {
		getLower().delete(txn, keys);
		for (Key key : keys) {
			cache.remove(key);
		}
	}

	@Override
	public Map<Key, Entity> get(Iterable<Key> keys) {
		Set<Key> others = new HashSet<Key>();
		Map<Key, Entity> result = new HashMap<Key, Entity>();

		for (Key key : keys) {
			Entity entity = cache.get(key);

			if (entity != null) {
				result.put(key, entity);
			} else {
				others.add(key);
			}

		}

		if (!others.isEmpty()) {
			Map<Key, Entity> tmp;

			result.putAll(tmp = getLower().get(others));
			cache.putAll(tmp);

		}

		return result;
	}

	@Override
	public Entity get(Key key) throws EntityNotFoundException {
		Entity entity;

		if ((entity = cache.get(key)) != null) {
			return entity;
		} else {
			entity = getLower().get(key);
			cache.put(key, entity);
			return entity;
		}

	}

	@Override
	public Map<Key, Entity> get(Key... keys) {
		return get(Arrays.asList(keys));
	}

	@Override
	public Map<Key, Entity> get(Transaction txn, Iterable<Key> keys) {
		Set<Key> others = new HashSet<Key>();
		Map<Key, Entity> result = new HashMap<Key, Entity>();

		for (Key key : keys) {
			Entity entity = cache.get(key);

			if (entity != null) {
				result.put(key, entity);
			} else {
				others.add(key);
			}

		}

		if (!others.isEmpty()) {
			Map<Key, Entity> tmp;

			result.putAll(tmp = getLower().get(txn, others));
			cache.putAll(tmp);

		}

		return result;
	}

	@Override
	public Entity get(Transaction txn, Key key) throws EntityNotFoundException {
		Entity entity;

		if ((entity = cache.get(key)) != null) {
			return entity;
		} else {
			entity = getLower().get(txn, key);
			cache.put(key, entity);
			return entity;
		}

	}

	@Override
	public Map<Key, Entity> get(Transaction txn, Key... keys) {
		return get(txn, Arrays.asList(keys));
	}

	@Override
	public void onPop(Catfish catfish) {

		getLower();
		getCatfish();

		this.lower = null;
		this.catfish = null;

	}

	@Override
	public void onPush(Catfish catfish, DataSource lower) {

		if (this.catfish != null || this.lower != null) {
			throw new IllegalStateException("Already bound to stack.");
		}

		this.lower = lower;
		this.catfish = catfish;

	}

	@Override
	public List<Key> put(Entity... entities) {
		return put(Arrays.asList(entities));
	}

	@Override
	public List<Key> put(Iterable<Entity> entities) {
		List<Key> keys = getLower().put(entities);
		Iterator<Key> keyIterator = keys.iterator();
		Iterator<Entity> entityIterator = entities.iterator();

		while (keyIterator.hasNext() && entityIterator.hasNext()) {
			cache.put(keyIterator.next(), entityIterator.next());
		}

		return keys;
	}

	@Override
	public List<Key> put(Transaction txn, Entity... entities) {
		return put(txn, Arrays.asList(entities));
	}

	@Override
	public List<Key> put(Transaction txn, Iterable<Entity> entities) {
		List<Key> keys = getLower().put(txn, entities);

		Iterator<Key> keyIterator = keys.iterator();
		Iterator<Entity> entityIterator = entities.iterator();

		while (keyIterator.hasNext() && entityIterator.hasNext()) {
			cache.put(keyIterator.next(), entityIterator.next());
		}

		return keys;
	}

	protected Catfish getCatfish() {

		if (catfish == null) {
			throw new IllegalStateException("Not bound to stack.");
		}

		return catfish;
	}

	protected DataSource getLower() {

		if (lower == null) {
			throw new IllegalStateException("Not bound to stack.");
		}

		return lower;
	}
}
